'''
合并单元格
确定权重展示表格

https://www.cnblogs.com/wutaotaosin/articles/12011167.html

import xlsxwriter

workbook = xlsxwriter.Workbook('hello.xlsx') # 建立文件

worksheet = workbook.add_worksheet() # 建立sheet， 可以work.add_worksheet('employee')来指定sheet名，但中文名会报UnicodeDecodeErro的错误

worksheet.write('A1', 'Hello world') # 向A1写入

workbook.close()

worksheet.merge_range('B2:E5', "", cell_format)

worksheet.write_rich_string('B2',
                            'This is ',
                            red, 'red',
                            ' and this is ',
                            blue, 'blue',
                            cell_format)
'''

from torch.nn.modules import batchnorm
from torchvision.models.inception import BasicConv2d
import xlsxwriter 
import os
import sys
import numpy as np
import matplotlib
from matplotlib import pyplot as plt
import torch
import math
import json
from torch.serialization import save

# 定义列表名称

def outConv1ToXls_matrix(xlspath,convlayer_dict):
    '''
    输出第一层卷积层的权重情况
    '''
    workbook = xlsxwriter.Workbook(xlspath)# 创建指定的表格
    for layer_name in convlayer_dict:
        # 以当前层确定工作目录
        worksheet=workbook.add_worksheet(layer_name)
        # 根据结构计算当前合并的情况
        # 主要计算整体的合并结果
        conv_medule=convlayer_dict[layer_name]['conv']
        batchnorm_medule=convlayer_dict[layer_name]['batchnorm']
        relu_medule=convlayer_dict[layer_name]['ReLU']
        # 确定具体参数
        conv_weight=conv_medule.weight
        batchnorm_bias=batchnorm_medule.bias
        batchnorm_gamma=batchnorm_medule.weight
        batchnorm_mean=batchnorm_medule.running_mean
        batchnorm_var=batchnorm_medule.running_var
        out_planes,in_planes,kernel_h,kernel_w=conv_weight.shape

        # 计算当前列名可能性
        column_num=1+in_planes*kernel_w+4 # 获得当前列数
        column_names=[]
        for c in range(26):
            for i in range(27):
                for j in range(26):
                    if i==0 and c==0:
                        column_names.append("{}".format(chr(ord('A')+j)))
                    elif c==0:
                        column_names.append("{}{}".format(chr(ord('A')+i-1),chr(ord('A')+j)))
                    else:
                        column_names.append("{}{}{}".format(chr(ord('A')+c),chr(ord('A')+i-1),chr(ord('A')+j)))
            if len(column_names)>2*column_num:
                break
        # 根据设计原则
        r_index,c_index=1,0
        # 确定列名
        worksheet.write('{}{}'.format(column_names[c_index],r_index),'Feature Index') # 列名
        c_index=c_index+1
        for k_c in range(in_planes):
            worksheet.merge_range('{}{}:{}{}'.format(column_names[c_index],r_index,column_names[c_index+kernel_w-1],r_index),'channel {}'.format(str(k_c)))
            c_index=c_index+kernel_w
        worksheet.write("{}{}".format(column_names[c_index+0],r_index),'bias')
        worksheet.write("{}{}".format(column_names[c_index+1],r_index),'gamma')
        worksheet.write("{}{}".format(column_names[c_index+2],r_index),'mean')
        worksheet.write("{}{}".format(column_names[c_index+3],r_index),'var')
        r_index=r_index+1
        for i in range(out_planes):
            c_index=0
            worksheet.merge_range('{}{}:{}{}'.format(column_names[c_index],r_index,column_names[c_index],r_index+kernel_h-1),str(int(i))) # 确定当前通道情况
            c_index=c_index+1
            # 核权重 写入
            for k_c in range(in_planes):
                for k_h in range(kernel_h):
                    for k_w in range(kernel_w):
                        worksheet.write('{}{}'.format(column_names[c_index+k_w],r_index+k_h),conv_weight[i,k_c,k_h,k_w])
                c_index=c_index+kernel_w # 确定列号
            # 输出batchnorm 的参数
            worksheet.merge_range('{}{}:{}{}'.format(column_names[c_index+0],r_index,column_names[c_index+0],r_index+kernel_h-1),batchnorm_bias[i]) # bias
            worksheet.merge_range('{}{}:{}{}'.format(column_names[c_index+1],r_index,column_names[c_index+1],r_index+kernel_h-1),batchnorm_gamma[i]) # gamma
            worksheet.merge_range('{}{}:{}{}'.format(column_names[c_index+2],r_index,column_names[c_index+2],r_index+kernel_h-1),batchnorm_mean[i]) # mean
            worksheet.merge_range('{}{}:{}{}'.format(column_names[c_index+3],r_index,column_names[c_index+3],r_index+kernel_h-1),batchnorm_var[i]) # var
            r_index=r_index+kernel_h # 更新行 
    workbook.close()
    print("out xlxs:{}".format(xlspath))


def outConv1ToXls_Single(xlspath,convlayer_dict):
    '''
    输出第一层卷积层的权重情况
    '''
    workbook = xlsxwriter.Workbook(xlspath)# 创建指定的表格
    for layer_name in convlayer_dict:
        # 以当前层确定工作目录
        worksheet=workbook.add_worksheet(layer_name)
        # 根据结构计算当前合并的情况
        # 主要计算整体的合并结果
        conv_medule=convlayer_dict[layer_name]['conv']
        batchnorm_medule=convlayer_dict[layer_name]['batchnorm']
        relu_medule=convlayer_dict[layer_name]['ReLU']
        # 确定具体参数
        conv_weight=conv_medule.weight
        batchnorm_bias=batchnorm_medule.bias
        batchnorm_gamma=batchnorm_medule.weight
        batchnorm_mean=batchnorm_medule.running_mean
        batchnorm_var=batchnorm_medule.running_var
        out_planes,in_planes,kernel_h,kernel_w=conv_weight.shape

        # 计算当前列名可能性
        column_num=1+in_planes*kernel_w*kernel_h+4 # 获得当前列数
        column_names=[]
        for c in range(26):
            for i in range(27):
                for j in range(26):
                    if i==0 and c==0:
                        column_names.append("{}".format(chr(ord('A')+j)))
                    elif c==0:
                        column_names.append("{}{}".format(chr(ord('A')+i-1),chr(ord('A')+j)))
                    else:
                        column_names.append("{}{}{}".format(chr(ord('A')+c),chr(ord('A')+i-1),chr(ord('A')+j)))
            if len(column_names)>2*column_num:
                break
        # 根据设计原则
        r_index,c_index=1,0
        # 确定列名
        worksheet.write('{}{}'.format(column_names[c_index],r_index),'Feature Index') # 列名
        c_index=c_index+1
        for k_c in range(in_planes):
            worksheet.merge_range('{}{}:{}{}'.format(column_names[c_index],r_index,column_names[c_index+kernel_w*kernel_h-1],r_index),'channel {}'.format(str(k_c)))
            c_index=c_index+kernel_w*kernel_h
        worksheet.write("{}{}".format(column_names[c_index+0],r_index),'bias')
        worksheet.write("{}{}".format(column_names[c_index+1],r_index),'gamma')
        worksheet.write("{}{}".format(column_names[c_index+2],r_index),'mean')
        worksheet.write("{}{}".format(column_names[c_index+3],r_index),'var')
        r_index=r_index+1
        for i in range(out_planes):
            c_index=0
            worksheet.write('{}{}'.format(column_names[c_index],r_index),str(int(i))) # 确定当前通道情况
            c_index=c_index+1
            # 核权重 写入
            for k_c in range(in_planes):
                for k_h in range(kernel_h):
                    for k_w in range(kernel_w):
                        worksheet.write('{}{}'.format(column_names[c_index+kernel_w*k_h+k_w],r_index),conv_weight[i,k_c,k_h,k_w])
                c_index=c_index+kernel_w*kernel_h # 确定列号
            # 输出batchnorm 的参数
            worksheet.write('{}{}'.format(column_names[c_index+0],r_index),batchnorm_bias[i]) # bias
            worksheet.write('{}{}'.format(column_names[c_index+1],r_index),batchnorm_gamma[i]) # gamma
            worksheet.write('{}{}'.format(column_names[c_index+2],r_index),batchnorm_mean[i]) # mean
            worksheet.write('{}{}'.format(column_names[c_index+3],r_index),batchnorm_var[i]) # var
            r_index=r_index+1 # 更新行 
    workbook.close()
    print("out xlxs:{}".format(xlspath))
         
